package cn.doourbest.mybatis.generate;

import org.mybatis.generator.api.JavaFormatter;
import org.mybatis.generator.api.dom.java.CompilationUnit;
import org.mybatis.generator.api.dom.java.FullyQualifiedJavaType;
import org.mybatis.generator.config.Context;

import java.util.Set;

/**
 * This class is the default formatter for generated Java.  This class will use the
 * built in formatting of the DOM classes directly.
 * 
 * @author Jeff Butler
 *
 */
public class CustomerDefaultJavaFormatter implements JavaFormatter {
    protected Context context;

    @Override
    public String getFormattedContent(CompilationUnit compilationUnit) {
        if(compilationUnit.getImportedTypes().size() == 1){
            for (FullyQualifiedJavaType importedType : compilationUnit.getImportedTypes()) {
                if (importedType.getShortName().contains("BaseMapper")){
                    FullyQualifiedJavaType fullyQualifiedJavaType = new FullyQualifiedJavaType(importedType.getFullyQualifiedNameWithoutTypeParameters());
                    compilationUnit.addImportedType(fullyQualifiedJavaType);

                    FullyQualifiedJavaType fullyQualifiedJavaType1 = new FullyQualifiedJavaType(importedType.getTypeArguments().get(0).getFullyQualifiedNameWithoutTypeParameters());
                    compilationUnit.addImportedType(fullyQualifiedJavaType1);
                }
            }
        }


        return compilationUnit.getFormattedContent();
    }

    @Override
    public void setContext(Context context) {
        this.context = context;
    }
}
